{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "Growing Neural Cellular Automata",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "28S76DVlfCMZ"
      },
      "source": [
        "# Growing Neural Cellular Automata\n",
        "\n",
        "This notebook contains code to reproduce experiments and figures for the [\"Growing Neural Cellular Automata\"](http://distill.pub/2020/growing-ca) article.\n",
        "\n",
        "Copyright 2020 Google LLC\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at\n",
        "\n",
        "    https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "id": "i5wi_r4gyzFr"
      },
      "source": [
        "#@title Imports and Notebook Utilities\n",
        "%tensorflow_version 2.x\n",
        "\n",
        "import os\n",
        "import io\n",
        "import PIL.Image, PIL.ImageDraw\n",
        "import base64\n",
        "import zipfile\n",
        "import json\n",
        "import requests\n",
        "import numpy as np\n",
        "import matplotlib.pylab as pl\n",
        "import glob\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "from IPython.display import Image, HTML, clear_output\n",
        "import tqdm\n",
        "\n",
        "import os\n",
        "os.environ['FFMPEG_BINARY'] = 'ffmpeg'\n",
        "import moviepy.editor as mvp\n",
        "from moviepy.video.io.ffmpeg_writer import FFMPEG_VideoWriter\n",
        "clear_output()\n",
        "\n",
        "def np2pil(a):\n",
        "  if a.dtype in [np.float32, np.float64]:\n",
        "    a = np.uint8(np.clip(a, 0, 1)*255)\n",
        "  return PIL.Image.fromarray(a)\n",
        "\n",
        "def imwrite(f, a, fmt=None):\n",
        "  a = np.asarray(a)\n",
        "  if isinstance(f, str):\n",
        "    fmt = f.rsplit('.', 1)[-1].lower()\n",
        "    if fmt == 'jpg':\n",
        "      fmt = 'jpeg'\n",
        "    f = open(f, 'wb')\n",
        "  np2pil(a).save(f, fmt, quality=95)\n",
        "\n",
        "def imencode(a, fmt='jpeg'):\n",
        "  a = np.asarray(a)\n",
        "  if len(a.shape) == 3 and a.shape[-1] == 4:\n",
        "    fmt = 'png'\n",
        "  f = io.BytesIO()\n",
        "  imwrite(f, a, fmt)\n",
        "  return f.getvalue()\n",
        "\n",
        "def im2url(a, fmt='jpeg'):\n",
        "  encoded = imencode(a, fmt)\n",
        "  base64_byte_string = base64.b64encode(encoded).decode('ascii')\n",
        "  return 'data:image/' + fmt.upper() + ';base64,' + base64_byte_string\n",
        "\n",
        "def imshow(a, fmt='jpeg'):\n",
        "  display(Image(data=imencode(a, fmt)))\n",
        "\n",
        "def tile2d(a, w=None):\n",
        "  a = np.asarray(a)\n",
        "  if w is None:\n",
        "    w = int(np.ceil(np.sqrt(len(a))))\n",
        "  th, tw = a.shape[1:3]\n",
        "  pad = (w-len(a))%w\n",
        "  a = np.pad(a, [(0, pad)]+[(0, 0)]*(a.ndim-1), 'constant')\n",
        "  h = len(a)//w\n",
        "  a = a.reshape([h, w]+list(a.shape[1:]))\n",
        "  a = np.rollaxis(a, 2, 1).reshape([th*h, tw*w]+list(a.shape[4:]))\n",
        "  return a\n",
        "\n",
        "def zoom(img, scale=4):\n",
        "  img = np.repeat(img, scale, 0)\n",
        "  img = np.repeat(img, scale, 1)\n",
        "  return img\n",
        "\n",
        "class VideoWriter:\n",
        "  def __init__(self, filename, fps=30.0, **kw):\n",
        "    self.writer = None\n",
        "    self.params = dict(filename=filename, fps=fps, **kw)\n",
        "\n",
        "  def add(self, img):\n",
        "    img = np.asarray(img)\n",
        "    if self.writer is None:\n",
        "      h, w = img.shape[:2]\n",
        "      self.writer = FFMPEG_VideoWriter(size=(w, h), **self.params)\n",
        "    if img.dtype in [np.float32, np.float64]:\n",
        "      img = np.uint8(img.clip(0, 1)*255)\n",
        "    if len(img.shape) == 2:\n",
        "      img = np.repeat(img[..., None], 3, -1)\n",
        "    self.writer.write_frame(img)\n",
        "\n",
        "  def close(self):\n",
        "    if self.writer:\n",
        "      self.writer.close()\n",
        "\n",
        "  def __enter__(self):\n",
        "    return self\n",
        "\n",
        "  def __exit__(self, *kw):\n",
        "    self.close()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "id": "zR6I1JONmWBb"
      },
      "source": [
        "#@title Cellular Automata Parameters\n",
        "CHANNEL_N = 16        # Number of CA state channels\n",
        "TARGET_PADDING = 16   # Number of pixels used to pad the target image border\n",
        "TARGET_SIZE = 40\n",
        "BATCH_SIZE = 8\n",
        "POOL_SIZE = 1024\n",
        "CELL_FIRE_RATE = 0.5\n",
        "\n",
        "TARGET_EMOJI = \"🦎\" #@param {type:\"string\"}\n",
        "\n",
        "EXPERIMENT_TYPE = \"Regenerating\" #@param [\"Growing\", \"Persistent\", \"Regenerating\"]\n",
        "EXPERIMENT_MAP = {\"Growing\":0, \"Persistent\":1, \"Regenerating\":2}\n",
        "EXPERIMENT_N = EXPERIMENT_MAP[EXPERIMENT_TYPE]\n",
        "\n",
        "USE_PATTERN_POOL = [0, 1, 1][EXPERIMENT_N]\n",
        "DAMAGE_N = [0, 0, 3][EXPERIMENT_N]  # Number of patterns to damage in a batch"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "id": "lCbPFbI_zosW"
      },
      "source": [
        "#@title CA Model and Utilities\n",
        "\n",
        "from tensorflow.keras.layers import Conv2D\n",
        "\n",
        "def load_image(url, max_size=TARGET_SIZE):\n",
        "  r = requests.get(url)\n",
        "  img = PIL.Image.open(io.BytesIO(r.content))\n",
        "  img.thumbnail((max_size, max_size), PIL.Image.ANTIALIAS)\n",
        "  img = np.float32(img)/255.0\n",
        "  # premultiply RGB by Alpha\n",
        "  img[..., :3] *= img[..., 3:]\n",
        "  return img\n",
        "\n",
        "def load_emoji(emoji):\n",
        "  code = hex(ord(emoji))[2:].lower()\n",
        "  url = 'https://github.com/googlefonts/noto-emoji/blob/main/png/128/emoji_u%s.png?raw=true'%code\n",
        "  return load_image(url)\n",
        "\n",
        "\n",
        "def to_rgba(x):\n",
        "  return x[..., :4]\n",
        "\n",
        "def to_alpha(x):\n",
        "  return tf.clip_by_value(x[..., 3:4], 0.0, 1.0)\n",
        "\n",
        "def to_rgb(x):\n",
        "  # assume rgb premultiplied by alpha\n",
        "  rgb, a = x[..., :3], to_alpha(x)\n",
        "  return 1.0-a+rgb\n",
        "\n",
        "def get_living_mask(x):\n",
        "  alpha = x[:, :, :, 3:4]\n",
        "  return tf.nn.max_pool2d(alpha, 3, [1, 1, 1, 1], 'SAME') > 0.1\n",
        "\n",
        "def make_seed(size, n=1):\n",
        "  x = np.zeros([n, size, size, CHANNEL_N], np.float32)\n",
        "  x[:, size//2, size//2, 3:] = 1.0\n",
        "  return x\n",
        "\n",
        "\n",
        "class CAModel(tf.keras.Model):\n",
        "\n",
        "  def __init__(self, channel_n=CHANNEL_N, fire_rate=CELL_FIRE_RATE):\n",
        "    super().__init__()\n",
        "    self.channel_n = channel_n\n",
        "    self.fire_rate = fire_rate\n",
        "\n",
        "    self.dmodel = tf.keras.Sequential([\n",
        "          Conv2D(128, 1, activation=tf.nn.relu),\n",
        "          Conv2D(self.channel_n, 1, activation=None,\n",
        "              kernel_initializer=tf.zeros_initializer),\n",
        "    ])\n",
        "\n",
        "    self(tf.zeros([1, 3, 3, channel_n]))  # dummy call to build the model\n",
        "\n",
        "  @tf.function\n",
        "  def perceive(self, x, angle=0.0):\n",
        "    identify = np.float32([0, 1, 0])\n",
        "    identify = np.outer(identify, identify)\n",
        "    dx = np.outer([1, 2, 1], [-1, 0, 1]) / 8.0  # Sobel filter\n",
        "    dy = dx.T\n",
        "    c, s = tf.cos(angle), tf.sin(angle)\n",
        "    kernel = tf.stack([identify, c*dx-s*dy, s*dx+c*dy], -1)[:, :, None, :]\n",
        "    kernel = tf.repeat(kernel, self.channel_n, 2)\n",
        "    y = tf.nn.depthwise_conv2d(x, kernel, [1, 1, 1, 1], 'SAME')\n",
        "    return y\n",
        "\n",
        "  @tf.function\n",
        "  def call(self, x, fire_rate=None, angle=0.0, step_size=1.0):\n",
        "    pre_life_mask = get_living_mask(x)\n",
        "\n",
        "    y = self.perceive(x, angle)\n",
        "    dx = self.dmodel(y)*step_size\n",
        "    if fire_rate is None:\n",
        "      fire_rate = self.fire_rate\n",
        "    update_mask = tf.random.uniform(tf.shape(x[:, :, :, :1])) <= fire_rate\n",
        "    x += dx * tf.cast(update_mask, tf.float32)\n",
        "\n",
        "    post_life_mask = get_living_mask(x)\n",
        "    life_mask = pre_life_mask & post_life_mask\n",
        "    return x * tf.cast(life_mask, tf.float32)\n",
        "\n",
        "\n",
        "CAModel().dmodel.summary()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xDX5HL7VLd0z"
      },
      "source": [
        "# Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "id": "IeWf6HeTe8kM"
      },
      "source": [
        "#@title Train Utilities (SamplePool, Model Export, Damage)\n",
        "from google.protobuf.json_format import MessageToDict\n",
        "from tensorflow.python.framework import convert_to_constants\n",
        "\n",
        "class SamplePool:\n",
        "  def __init__(self, *, _parent=None, _parent_idx=None, **slots):\n",
        "    self._parent = _parent\n",
        "    self._parent_idx = _parent_idx\n",
        "    self._slot_names = slots.keys()\n",
        "    self._size = None\n",
        "    for k, v in slots.items():\n",
        "      if self._size is None:\n",
        "        self._size = len(v)\n",
        "      assert self._size == len(v)\n",
        "      setattr(self, k, np.asarray(v))\n",
        "\n",
        "  def sample(self, n):\n",
        "    idx = np.random.choice(self._size, n, False)\n",
        "    batch = {k: getattr(self, k)[idx] for k in self._slot_names}\n",
        "    batch = SamplePool(**batch, _parent=self, _parent_idx=idx)\n",
        "    return batch\n",
        "\n",
        "  def commit(self):\n",
        "    for k in self._slot_names:\n",
        "      getattr(self._parent, k)[self._parent_idx] = getattr(self, k)\n",
        "\n",
        "@tf.function\n",
        "def make_circle_masks(n, h, w):\n",
        "  x = tf.linspace(-1.0, 1.0, w)[None, None, :]\n",
        "  y = tf.linspace(-1.0, 1.0, h)[None, :, None]\n",
        "  center = tf.random.uniform([2, n, 1, 1], -0.5, 0.5)\n",
        "  r = tf.random.uniform([n, 1, 1], 0.1, 0.4)\n",
        "  x, y = (x-center[0])/r, (y-center[1])/r\n",
        "  mask = tf.cast(x*x+y*y < 1.0, tf.float32)\n",
        "  return mask\n",
        "\n",
        "def export_model(ca, base_fn):\n",
        "  ca.save_weights(base_fn)\n",
        "\n",
        "  cf = ca.call.get_concrete_function(\n",
        "      x=tf.TensorSpec([None, None, None, CHANNEL_N]),\n",
        "      fire_rate=tf.constant(0.5),\n",
        "      angle=tf.constant(0.0),\n",
        "      step_size=tf.constant(1.0))\n",
        "  cf = convert_to_constants.convert_variables_to_constants_v2(cf)\n",
        "  graph_def = cf.graph.as_graph_def()\n",
        "  graph_json = MessageToDict(graph_def)\n",
        "  graph_json['versions'] = dict(producer='1.14', minConsumer='1.14')\n",
        "  model_json = {\n",
        "      'format': 'graph-model',\n",
        "      'modelTopology': graph_json,\n",
        "      'weightsManifest': [],\n",
        "  }\n",
        "  with open(base_fn+'.json', 'w') as f:\n",
        "    json.dump(model_json, f)\n",
        "\n",
        "def generate_pool_figures(pool, step_i):\n",
        "  tiled_pool = tile2d(to_rgb(pool.x[:49]))\n",
        "  fade = np.linspace(1.0, 0.0, 72)\n",
        "  ones = np.ones(72) \n",
        "  tiled_pool[:, :72] += (-tiled_pool[:, :72] + ones[None, :, None]) * fade[None, :, None] \n",
        "  tiled_pool[:, -72:] += (-tiled_pool[:, -72:] + ones[None, :, None]) * fade[None, ::-1, None]\n",
        "  tiled_pool[:72, :] += (-tiled_pool[:72, :] + ones[:, None, None]) * fade[:, None, None]\n",
        "  tiled_pool[-72:, :] += (-tiled_pool[-72:, :] + ones[:, None, None]) * fade[::-1, None, None]\n",
        "  imwrite('train_log/%04d_pool.jpg'%step_i, tiled_pool)\n",
        "\n",
        "def visualize_batch(x0, x, step_i):\n",
        "  vis0 = np.hstack(to_rgb(x0).numpy())\n",
        "  vis1 = np.hstack(to_rgb(x).numpy())\n",
        "  vis = np.vstack([vis0, vis1])\n",
        "  imwrite('train_log/batches_%04d.jpg'%step_i, vis)\n",
        "  print('batch (before/after):')\n",
        "  imshow(vis)\n",
        "\n",
        "def plot_loss(loss_log):\n",
        "  pl.figure(figsize=(10, 4))\n",
        "  pl.title('Loss history (log10)')\n",
        "  pl.plot(np.log10(loss_log), '.', alpha=0.1)\n",
        "  pl.show()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "id": "zKlA50h0jlvl"
      },
      "source": [
        "#@title Choose Target Image { vertical-output: true}\n",
        "#url = 'https://github.com/google-research/self-organising-systems/blob/master/assets/growing_ca/planaria2_48.png?raw=true'\n",
        "#target_img = load_image(url, 48)\n",
        "\n",
        "target_img = load_emoji(TARGET_EMOJI)\n",
        "imshow(zoom(to_rgb(target_img), 2), fmt='png')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "id": "ak5rBmbxmHV7"
      },
      "source": [
        "#@title Initialize Training { vertical-output: true}\n",
        "\n",
        "p = TARGET_PADDING\n",
        "pad_target = tf.pad(target_img, [(p, p), (p, p), (0, 0)])\n",
        "h, w = pad_target.shape[:2]\n",
        "seed = np.zeros([h, w, CHANNEL_N], np.float32)\n",
        "seed[h//2, w//2, 3:] = 1.0\n",
        "\n",
        "def loss_f(x):\n",
        "  return tf.reduce_mean(tf.square(to_rgba(x)-pad_target), [-2, -3, -1])\n",
        "\n",
        "ca = CAModel()\n",
        "\n",
        "loss_log = []\n",
        "\n",
        "lr = 2e-3\n",
        "lr_sched = tf.keras.optimizers.schedules.PiecewiseConstantDecay(\n",
        "    [2000], [lr, lr*0.1])\n",
        "trainer = tf.keras.optimizers.Adam(lr_sched)\n",
        "\n",
        "loss0 = loss_f(seed).numpy()\n",
        "pool = SamplePool(x=np.repeat(seed[None, ...], POOL_SIZE, 0))\n",
        "\n",
        "!mkdir -p train_log && rm -f train_log/*"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "id": "QzP_vDchq0d9"
      },
      "source": [
        "#@title Training Loop {vertical-output: true}\n",
        "\n",
        "@tf.function\n",
        "def train_step(x):\n",
        "  iter_n = tf.random.uniform([], 64, 96, tf.int32)\n",
        "  with tf.GradientTape() as g:\n",
        "    for i in tf.range(iter_n):\n",
        "      x = ca(x)\n",
        "    loss = tf.reduce_mean(loss_f(x))\n",
        "  grads = g.gradient(loss, ca.weights)\n",
        "  grads = [g/(tf.norm(g)+1e-8) for g in grads]\n",
        "  trainer.apply_gradients(zip(grads, ca.weights))\n",
        "  return x, loss\n",
        "\n",
        "for i in range(8000+1):\n",
        "  if USE_PATTERN_POOL:\n",
        "    batch = pool.sample(BATCH_SIZE)\n",
        "    x0 = batch.x\n",
        "    loss_rank = loss_f(x0).numpy().argsort()[::-1]\n",
        "    x0 = x0[loss_rank]\n",
        "    x0[:1] = seed\n",
        "    if DAMAGE_N:\n",
        "      damage = 1.0-make_circle_masks(DAMAGE_N, h, w).numpy()[..., None]\n",
        "      x0[-DAMAGE_N:] *= damage\n",
        "  else:\n",
        "    x0 = np.repeat(seed[None, ...], BATCH_SIZE, 0)\n",
        "\n",
        "  x, loss = train_step(x0)\n",
        "\n",
        "  if USE_PATTERN_POOL:\n",
        "    batch.x[:] = x\n",
        "    batch.commit()\n",
        "\n",
        "  step_i = len(loss_log)\n",
        "  loss_log.append(loss.numpy())\n",
        "  \n",
        "  if step_i%10 == 0:\n",
        "    generate_pool_figures(pool, step_i)\n",
        "  if step_i%100 == 0:\n",
        "    clear_output()\n",
        "    visualize_batch(x0, x, step_i)\n",
        "    plot_loss(loss_log)\n",
        "    export_model(ca, 'train_log/%04d'%step_i)\n",
        "\n",
        "  print('\\r step: %d, log10(loss): %.3f'%(len(loss_log), np.log10(loss)), end='')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QAscSKkRaFwp"
      },
      "source": [
        "# Figures"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "id": "nqvkfl9W4ODI"
      },
      "source": [
        "#@title Training Progress (Checkpoints)\n",
        "\n",
        "models = []\n",
        "for i in [100, 500, 1000, 4000]:\n",
        "  ca = CAModel()\n",
        "  ca.load_weights('train_log/%04d'%i)\n",
        "  models.append(ca)\n",
        "\n",
        "out_fn = 'train_steps_damage_%d.mp4'%DAMAGE_N\n",
        "x = np.zeros([len(models), 72, 72, CHANNEL_N], np.float32)\n",
        "x[..., 36, 36, 3:] = 1.0\n",
        "with VideoWriter(out_fn) as vid:\n",
        "  for i in tqdm.trange(500):\n",
        "    vis = np.hstack(to_rgb(x))\n",
        "    vid.add(zoom(vis, 2))\n",
        "    for ca, xk in zip(models, x):\n",
        "      xk[:] = ca(xk[None,...])[0]\n",
        "mvp.ipython_display(out_fn)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "id": "QeXZKb5v2gxj"
      },
      "source": [
        "#@title Training Progress (Batches)\n",
        "frames = sorted(glob.glob('train_log/batches_*.jpg'))\n",
        "mvp.ImageSequenceClip(frames, fps=10.0).write_videofile('batches.mp4')\n",
        "mvp.ipython_display('batches.mp4')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "id": "B4JAbAJf6Alw"
      },
      "source": [
        "#@title Pool Contents\n",
        "frames = sorted(glob.glob('train_log/*_pool.jpg'))[:80]\n",
        "mvp.ImageSequenceClip(frames, fps=20.0).write_videofile('pool.mp4')\n",
        "mvp.ipython_display('pool.mp4')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lyxeGm6dJX8D"
      },
      "source": [
        "## Pretrained Models and Figures\n",
        "\n",
        "Please run the cell below to download pretrained models that are used to generate the subsequent figures. The figures generated after this are generated using the pretrained CAs."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wiGl7S0E6-OA"
      },
      "source": [
        "!wget -O models.zip 'https://github.com/google-research/self-organising-systems/blob/master/assets/growing_ca/models.zip?raw=true'\n",
        "!unzip -oq models.zip\n",
        "\n",
        "EMOJI = '🦎😀💥👁🐠🦋🐞🕸🥨🎄'\n",
        "\n",
        "def get_model(emoji='🦋', fire_rate=0.5, use_pool=1, damage_n=3, run=0,\n",
        "              prefix='models/', output='model'):\n",
        "  path = prefix\n",
        "  assert fire_rate in [0.5, 1.0]\n",
        "  if fire_rate==0.5:\n",
        "    path += 'use_sample_pool_%d damage_n_%d '%(use_pool, damage_n)\n",
        "  elif fire_rate==1.0:\n",
        "    path += 'fire_rate_1.0 '\n",
        "  code = hex(ord(emoji))[2:].upper()\n",
        "  path += 'target_emoji_%s run_index_%d/08000'%(code, run)\n",
        "  assert output in ['model', 'json']\n",
        "  if output == 'model':\n",
        "    ca = CAModel(channel_n=16, fire_rate=fire_rate)\n",
        "    ca.load_weights(path)\n",
        "    return ca\n",
        "  elif output == 'json':\n",
        "    return open(path+'.json', 'r').read()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YyMms2wKwX9x"
      },
      "source": [
        "atlas = np.hstack([load_emoji(e) for e in EMOJI])\n",
        "imshow(atlas)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "id": "mqgtL5VpLEeL"
      },
      "source": [
        "#@title Teaser\n",
        "models = [get_model(emoji, run=1) for emoji in EMOJI]\n",
        "\n",
        "with VideoWriter('teaser.mp4') as vid:\n",
        "  x = np.zeros([len(EMOJI), 64, 64, CHANNEL_N], np.float32)\n",
        "  # grow\n",
        "  for i in tqdm.trange(200):\n",
        "    k = i//20\n",
        "    if i%20==0 and k<len(EMOJI):\n",
        "      x[k, 32, 32, 3:] = 1.0\n",
        "    vid.add(zoom(tile2d(to_rgb(x), 5), 2))\n",
        "    for ca, xk in zip(models, x):\n",
        "      xk[:] = ca(xk[None,...])[0]\n",
        "  # damage\n",
        "  mask = PIL.Image.new('L', (64*5, 64*2))\n",
        "  draw = PIL.ImageDraw.Draw(mask)\n",
        "  for i in tqdm.trange(400):\n",
        "    cx, r = i*3-20, 6\n",
        "    y1, y2 = 32+np.sin(i/5+np.pi)*8, 32+64+np.sin(i/5)*8\n",
        "    draw.rectangle((0, 0, 64*5, 64*2), fill=0)\n",
        "    draw.ellipse((cx-r, y1-r, cx+r, y1+r), fill=255)\n",
        "    draw.ellipse((cx-r, y2-r, cx+r, y2+r), fill=255)\n",
        "    x *= 1.0-(np.float32(mask).reshape(2, 64, 5, 64)\n",
        "        .transpose([0, 2, 1, 3]).reshape(10, 64, 64, 1))/255.0\n",
        "    if i<200 or i%2 == 0:\n",
        "      vid.add(zoom(tile2d(to_rgb(x), 5), 2))\n",
        "    for ca, xk in zip(models, x):\n",
        "      xk[:] = ca(xk[None,...])[0]\n",
        "  # fade out\n",
        "  last = zoom(tile2d(to_rgb(x), 5), 2)\n",
        "  for t in np.linspace(0, 1, 30):\n",
        "    vid.add(last*(1.0-t)+t)\n",
        "\n",
        "mvp.ipython_display('teaser.mp4', loop=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "id": "4O4tzfe-GRJ7"
      },
      "source": [
        "#@title Unstable Patterns\n",
        "!wget -O slider.png 'https://github.com/google-research/self-organising-systems/raw/master/assets/growing_ca/slider.png?raw=true'\n",
        "\n",
        "import PIL.ImageFont\n",
        "from matplotlib import font_manager as fm\n",
        "font_fn = fm.findfont(fm.FontProperties())\n",
        "font = PIL.ImageFont.truetype(font_fn, 20)\n",
        "\n",
        "models = [get_model(ch, use_pool=0, damage_n=0) for ch in EMOJI]\n",
        "fn = 'unstable.mp4'\n",
        "with VideoWriter(fn) as vid:\n",
        "  x = np.zeros([len(EMOJI), 64, 64, CHANNEL_N], np.float32)\n",
        "  x[:, 32, 32, 3:] = 1.0\n",
        "  # grow\n",
        "  slider = PIL.Image.open(\"slider.png\")\n",
        "  for i in tqdm.trange(1000):\n",
        "    if i<200 or i%5 == 0:\n",
        "      vis = zoom(tile2d(to_rgb(x), 5), 4).clip(0, 1)\n",
        "      vis_extended = np.concatenate((vis, np.ones((164, vis.shape[1], 3))), axis=0) \n",
        "      im = np.uint8(vis_extended*255)\n",
        "      im = PIL.Image.fromarray(im)\n",
        "      im.paste(slider, box=(20, vis.shape[0]+20))\n",
        "      draw = PIL.ImageDraw.Draw(im)\n",
        "      p_x = (14 + (610/1000)*i)*2.0\n",
        "      draw.rectangle([p_x, vis.shape[0]+20+55, p_x+10, vis.shape[0]+20+82], fill=\"#434343bd\")\n",
        "      vid.add(np.uint8(im))\n",
        "    for ca, xk in zip(models, x):\n",
        "      xk[:] = ca(xk[None,...])[0]\n",
        "  # fade out\n",
        "  for t in np.linspace(0, 1, 30):\n",
        "    vid.add(vis_extended*(1.0-t)+t)\n",
        "\n",
        "mvp.ipython_display(fn, loop=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "id": "1CVR9MeYnjuY"
      },
      "source": [
        "#@title Rotation\n",
        "row_size = 4\n",
        "models_of_interest = [\"🦋\",\"🦎\",\"🐠\",\"😀\"]\n",
        "num_images = 16\n",
        "imgs = []\n",
        "start_angle = np.random.randint(13, 76)\n",
        "\n",
        "for i in np.arange(num_images):\n",
        "  ang = start_angle + i * np.random.randint(36, 111)\n",
        "  ang = ang/360.0 * 2 * np.pi\n",
        "  if i % row_size == 0:\n",
        "    ca = get_model(models_of_interest[i // row_size])\n",
        "  x = np.zeros([1, 56, 56, CHANNEL_N], np.float32)\n",
        "  x[:, 28, 28, 3:] = 1.0\n",
        "  for i in range(500):\n",
        "    ang = tf.constant(ang, tf.float32)\n",
        "    x = ca(x, angle=ang)\n",
        "  imgs.append(to_rgb(x)[0])\n",
        "# Assumes the result is a multiple of row_size\n",
        "assert len(imgs) % row_size == 0\n",
        "imgs = zip(*(iter(imgs),) * row_size)\n",
        "\n",
        "imgs_arr = np.concatenate([np.hstack(im_row) for im_row in imgs])\n",
        "vis = zoom(imgs_arr, 4)\n",
        "\n",
        "imshow(vis, fmt='png')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "id": "S5JRLGxX1dnX"
      },
      "source": [
        "#@title Regeneration (trained without damage)\n",
        "models = [get_model(ch, damage_n=0) for ch in '😀🦋🦎']\n",
        "with VideoWriter('regen1.mp4') as vid:\n",
        "  x = np.zeros([len(models), 5, 56, 56, CHANNEL_N], np.float32)\n",
        "  cx, cy = 28, 28\n",
        "  x[:, :, cy, cx, 3:] = 1.0\n",
        "  for i in tqdm.trange(2000):\n",
        "    if i == 200:\n",
        "      x[:, 0, cy:] = x[:, 1, :cy] = 0\n",
        "      x[:, 2, :, cx:] = x[:, 3, :, :cx] = 0\n",
        "      x[:, 4, cy-8:cy+8, cx-8:cx+8] = 0\n",
        "    vis = to_rgb(x)\n",
        "    vis = np.vstack([np.hstack(row) for row in vis])\n",
        "    vis = zoom(vis, 2)\n",
        "    if (i < 400 and i%2==0) or i%8 == 0:\n",
        "      vid.add(vis)\n",
        "    if i == 200:\n",
        "      for _ in range(29):\n",
        "        vid.add(vis)\n",
        "    for ca, row in zip(models, x):\n",
        "      row[:] = ca(row)\n",
        "\n",
        "mvp.ipython_display('regen1.mp4')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "id": "TDzJM69u4_8p"
      },
      "source": [
        "#@title Regeneration (trained with damage)\n",
        "models = [get_model(ch, damage_n=3) for ch in '😀🦋🦎']\n",
        "with VideoWriter('regen2.mp4') as vid:\n",
        "  x = np.zeros([len(models), 5, 56, 56, CHANNEL_N], np.float32)\n",
        "  cx, cy = 28, 28\n",
        "  x[:, :, cy, cx, 3:] = 1.0\n",
        "  for i in tqdm.trange(2000):\n",
        "    if i == 200:\n",
        "      x[:, 0, cy:] = x[:, 1, :cy] = 0\n",
        "      x[:, 2, :, cx:] = x[:, 3, :, :cx] = 0\n",
        "      x[:, 4, cy-8:cy+8, cx-8:cx+8] = 0\n",
        "    vis = to_rgb(x)\n",
        "    vis = np.vstack([np.hstack(row) for row in vis])\n",
        "    vis = zoom(vis, 2)\n",
        "    if (i < 400 and i%2==0) or i%8 == 0:\n",
        "      vid.add(vis)\n",
        "    if i == 200:\n",
        "      for _ in range(29):\n",
        "        vid.add(vis)\n",
        "    for ca, row in zip(models, x):\n",
        "      row[:] = ca(row)\n",
        "\n",
        "mvp.ipython_display('regen2.mp4')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "id": "fQ1u2MqFy7Ni"
      },
      "source": [
        "#@title Planarian\n",
        "!wget -O planarian.zip 'https://github.com/google-research/self-organising-systems/blob/master/assets/growing_ca/planarian.zip?raw=true'\n",
        "!unzip -oq planarian.zip -d planarian\n",
        "\n",
        "ca = CAModel()\n",
        "ca.load_weights('planarian/train_log/8000')\n",
        "\n",
        "x = np.zeros([1, 64, 96, CHANNEL_N], np.float32)\n",
        "x[:, 32, 48, 3:] = 1.0\n",
        "with VideoWriter('planarian.mp4', 30.0) as vid:\n",
        "  for i in range(400):\n",
        "    vid.add(zoom(to_rgb(x[0])))\n",
        "    x = ca(x, angle=np.pi/2.0)\n",
        "    if i==150:\n",
        "      x = x.numpy()\n",
        "      for k in range(24):\n",
        "        x[:,:24] = np.roll(x[:,:24], 1, 2)\n",
        "        x[:,-24:] = np.roll(x[:,-24:], -1, 2)\n",
        "        vid.add(zoom(to_rgb(x[0])))\n",
        "      for k in range(20):\n",
        "        vid.add(zoom(to_rgb(x[0])))\n",
        "\n",
        "mvp.ipython_display('planarian.mp4')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9M-oDuhea7bR"
      },
      "source": [
        "# Interactive Demos\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "id": "D7ypa-b7_fTn"
      },
      "source": [
        "#@title TensorFlow.js Demo {run:\"auto\", vertical-output: true}\n",
        "#@markdown Select \"CHECKPOINT\" model to load the checkpoint created by running cells from the \"Training\" section of this notebook\n",
        "import IPython.display\n",
        "\n",
        "model = \"CHECKPOINT\"  #@param ['CHECKPOINT', '😀 1F600', '💥 1F4A5', '👁 1F441', '🦎 1F98E', '🐠 1F420', '🦋 1F98B', '🐞 1F41E', '🕸 1F578', '🥨 1F968', '🎄 1F384']\n",
        "model_type = '3 regenerating'  #@param ['1 naive', '2 persistent', '3 regenerating']\n",
        "\n",
        "#@markdown Shift-click to seed the pattern\n",
        "\n",
        "if model != 'CHECKPOINT':\n",
        "  code = model.split(' ')[1]\n",
        "  emoji = chr(int(code, 16))\n",
        "  experiment_i = int(model_type.split()[0])-1\n",
        "  use_pool = (0, 1, 1)[experiment_i]\n",
        "  damage_n = (0, 0, 3)[experiment_i]\n",
        "  model_str = get_model(emoji, use_pool=use_pool, damage_n=damage_n, output='json')\n",
        "else:\n",
        "  last_checkpoint_fn = sorted(glob.glob('train_log/*.json'))[-1]\n",
        "  model_str = open(last_checkpoint_fn).read()\n",
        "\n",
        "data_js = '''\n",
        "  window.GRAPH_URL = URL.createObjectURL(new Blob([`%s`], {type: 'application/json'}));\n",
        "'''%(model_str)\n",
        "\n",
        "display(IPython.display.Javascript(data_js))\n",
        "\n",
        "\n",
        "IPython.display.HTML('''\n",
        "<script src=\"https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@2.3.0/dist/tf.min.js\"></script>\n",
        "\n",
        "<canvas id='canvas' style=\"border: 1px solid black; image-rendering: pixelated;\"></canvas>\n",
        "\n",
        "<script>\n",
        "  \"use strict\";\n",
        "  \n",
        "  const sleep = (ms)=>new Promise(resolve => setTimeout(resolve, ms));\n",
        "  \n",
        "  const parseConsts = model_graph=>{\n",
        "    const dtypes = {'DT_INT32':['int32', 'intVal', Int32Array],\n",
        "                    'DT_FLOAT':['float32', 'floatVal', Float32Array]};\n",
        "    \n",
        "    const consts = {};\n",
        "    model_graph.modelTopology.node.filter(n=>n.op=='Const').forEach((node=>{\n",
        "      const v = node.attr.value.tensor;\n",
        "      const [dtype, field, arrayType] = dtypes[v.dtype];\n",
        "      if (!v.tensorShape.dim) {\n",
        "        consts[node.name] = [tf.scalar(v[field][0], dtype)];\n",
        "      } else {\n",
        "        // if there is a 0-length dimension, the exported graph json lacks \"size\"\n",
        "        const shape = v.tensorShape.dim.map(d=>(!d.size) ? 0 : parseInt(d.size));\n",
        "        let arr;\n",
        "        if (v.tensorContent) {\n",
        "          const data = atob(v.tensorContent);\n",
        "          const buf = new Uint8Array(data.length);\n",
        "          for (var i=0; i<data.length; ++i) {\n",
        "            buf[i] = data.charCodeAt(i);\n",
        "          }\n",
        "          arr = new arrayType(buf.buffer);\n",
        "        } else {\n",
        "          const size = shape.reduce((a, b)=>a*b);\n",
        "          arr = new arrayType(size);\n",
        "          if (size) {\n",
        "            arr.fill(v[field][0]);\n",
        "          }\n",
        "        }\n",
        "        consts[node.name] = [tf.tensor(arr, shape, dtype)];\n",
        "      }\n",
        "    }));\n",
        "    return consts;\n",
        "  }\n",
        "  \n",
        "  const run = async ()=>{\n",
        "    const r = await fetch(GRAPH_URL);\n",
        "    const consts = parseConsts(await r.json());\n",
        "    \n",
        "    const model = await tf.loadGraphModel(GRAPH_URL);\n",
        "    Object.assign(model.weights, consts);\n",
        "    \n",
        "    let seed = new Array(16).fill(0).map((x, i)=>i<3?0:1);\n",
        "    seed = tf.tensor(seed, [1, 1, 1, 16]);\n",
        "    \n",
        "    const D = 96;\n",
        "    const initState = tf.tidy(()=>{\n",
        "      const D2 = D/2;\n",
        "      const a = seed.pad([[0, 0], [D2-1, D2], [D2-1, D2], [0,0]]);\n",
        "      return a;\n",
        "    });\n",
        "    \n",
        "    const state = tf.variable(initState);\n",
        "    const [_, h, w, ch] = state.shape;\n",
        "    \n",
        "    const damage = (x, y, r)=>{\n",
        "      tf.tidy(()=>{\n",
        "        const rx = tf.range(0, w).sub(x).div(r).square().expandDims(0);\n",
        "        const ry = tf.range(0, h).sub(y).div(r).square().expandDims(1);\n",
        "        const mask = rx.add(ry).greater(1.0).expandDims(2);\n",
        "        state.assign(state.mul(mask));\n",
        "      });\n",
        "    }\n",
        "    \n",
        "    const plantSeed = (x, y)=>{\n",
        "      const x2 = w-x-seed.shape[2];\n",
        "      const y2 = h-y-seed.shape[1];\n",
        "      if (x<0 || x2<0 || y2<0 || y2<0)\n",
        "        return;\n",
        "      tf.tidy(()=>{\n",
        "        const a = seed.pad([[0, 0], [y, y2], [x, x2], [0,0]]);\n",
        "        state.assign(state.add(a));\n",
        "      });\n",
        "    }\n",
        "    \n",
        "    const scale = 4;\n",
        "    \n",
        "    const canvas = document.getElementById('canvas');\n",
        "    const ctx = canvas.getContext('2d');\n",
        "    canvas.width = w;\n",
        "    canvas.height = h;\n",
        "    canvas.style.width = `${w*scale}px`;\n",
        "    canvas.style.height = `${h*scale}px`;\n",
        "    \n",
        "    canvas.onmousedown = e=>{\n",
        "      const x = Math.floor(e.clientX/scale);\n",
        "        const y = Math.floor(e.clientY/scale);\n",
        "        if (e.buttons == 1) {\n",
        "          if (e.shiftKey) {\n",
        "            plantSeed(x, y);  \n",
        "          } else {\n",
        "            damage(x, y, 8);\n",
        "          }\n",
        "        }\n",
        "    }\n",
        "    canvas.onmousemove = e=>{\n",
        "      const x = Math.floor(e.clientX/scale);\n",
        "      const y = Math.floor(e.clientY/scale);\n",
        "      if (e.buttons == 1 && !e.shiftKey) {\n",
        "        damage(x, y, 8);\n",
        "      }\n",
        "    }\n",
        "\n",
        "    function step() {\n",
        "      tf.tidy(()=>{\n",
        "        state.assign(model.execute(\n",
        "            {x:state, fire_rate:tf.tensor(0.5),\n",
        "            angle:tf.tensor(0.0), step_size:tf.tensor(1.0)}, ['Identity']));\n",
        "      });\n",
        "    }\n",
        "\n",
        "    function render() {\n",
        "      step();\n",
        "\n",
        "      const imageData = tf.tidy(()=>{\n",
        "        const rgba = state.slice([0, 0, 0, 0], [-1, -1, -1, 4]);\n",
        "        const a = state.slice([0, 0, 0, 3], [-1, -1, -1, 1]);\n",
        "        const img = tf.tensor(1.0).sub(a).add(rgba).mul(255);\n",
        "        const rgbaBytes = new Uint8ClampedArray(img.dataSync());\n",
        "        return new ImageData(rgbaBytes, w, h);\n",
        "      });\n",
        "      ctx.putImageData(imageData, 0, 0);\n",
        "\n",
        "      requestAnimationFrame(render);\n",
        "    }\n",
        "    render();\n",
        "  }\n",
        "  run();\n",
        "  \n",
        "</script>\n",
        "''')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "id": "POma99rMIfV4"
      },
      "source": [
        "#@title WebGL Demo\n",
        "\n",
        "#@markdown This code exports quantized models for the WebGL demo that is used in the article.\n",
        "#@markdown The demo code can be found at https://github.com/distillpub/post--growing-ca/blob/master/public/ca.js\n",
        "\n",
        "def pack_layer(weight, bias, outputType=np.uint8):\n",
        "  in_ch, out_ch = weight.shape\n",
        "  assert (in_ch%4==0) and (out_ch%4==0) and (bias.shape==(out_ch,))\n",
        "  weight_scale, bias_scale = 1.0, 1.0\n",
        "  if outputType == np.uint8:\n",
        "    weight_scale = 2.0*np.abs(weight).max()\n",
        "    bias_scale = 2.0*np.abs(bias).max()\n",
        "    weight = np.round((weight/weight_scale+0.5)*255)\n",
        "    bias = np.round((bias/bias_scale+0.5)*255)\n",
        "  packed = np.vstack([weight, bias[None,...]])\n",
        "  packed = packed.reshape(in_ch+1, out_ch//4, 4)\n",
        "  packed = outputType(packed)\n",
        "  packed_b64 = base64.b64encode(packed.tobytes()).decode('ascii')\n",
        "  return {'data_b64': packed_b64, 'in_ch': in_ch, 'out_ch': out_ch,\n",
        "          'weight_scale': weight_scale, 'bias_scale': bias_scale,\n",
        "          'type': outputType.__name__}\n",
        "\n",
        "def export_ca_to_webgl_demo(ca, outputType=np.uint8):\n",
        "  # reorder the first layer inputs to meet webgl demo perception layout\n",
        "  chn = ca.channel_n\n",
        "  w1 = ca.weights[0][0, 0].numpy()\n",
        "  w1 = w1.reshape(chn, 3, -1).transpose(1, 0, 2).reshape(3*chn, -1)\n",
        "  layers = [\n",
        "      pack_layer(w1, ca.weights[1].numpy(), outputType),\n",
        "      pack_layer(ca.weights[2][0, 0].numpy(), ca.weights[3].numpy(), outputType)\n",
        "  ]\n",
        "  return json.dumps(layers)\n",
        "\n",
        "with zipfile.ZipFile('webgl_models8.zip', 'w') as zf:\n",
        "  for e in EMOJI:\n",
        "    zf.writestr('ex1_%s.json'%e, export_ca_to_webgl_demo(get_model(e, use_pool=0, damage_n=0)))\n",
        "    run = 1 if e in '😀🕸' else 0  # select runs that happen to quantize better\n",
        "    zf.writestr('ex2_%s.json'%e, export_ca_to_webgl_demo(get_model(e, use_pool=1, damage_n=0, run=run)))\n",
        "    run = 1 if e in '🦎' else 0    # select runs that happen to quantize better\n",
        "    zf.writestr('ex3_%s.json'%e, export_ca_to_webgl_demo(get_model(e, use_pool=1, damage_n=3, run=run)))"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}